package mx.com.crowdgine.util.gpucompute.jocl;

import java.io.IOException;
import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.HashMap;

import javax.media.opengl.GL2;
import javax.media.opengl.GLAutoDrawable;

import com.jogamp.opencl.CLBuffer;
import com.jogamp.opencl.CLCommandQueue;
import com.jogamp.opencl.CLDevice;
import com.jogamp.opencl.CLMemory;
import com.jogamp.opencl.CLDevice.Type;
import com.jogamp.opencl.CLKernel;
import com.jogamp.opencl.CLPlatform;
import com.jogamp.opencl.CLProgram;
import com.jogamp.opencl.CLContext;
import com.jogamp.opencl.gl.CLGLBuffer;
import com.jogamp.opencl.gl.CLGLContext;
import com.jogamp.opencl.CL;


public class JoclGpuCompute {
	public static final int CROWD_OCL_CPU_ONLY = 0;
	public static final int CROWD_OCL_GPU_ONLY = 1;
	public static final int CROWD_OCL_BEST_AVAILABLE = 2;
	
	private CLGLContext clGLContext;
	private CLContext clContext;
	private boolean usingGLInterop = false;
	private boolean verbose = false;
	
	//Public for testing, change later
	public HashMap<String, CLKernel> kernels;
	public CLCommandQueue commandQueue;
	public HashMap<String, CLBuffer<FloatBuffer>> clBuffers;
	public HashMap<String, CLGLBuffer<?>> clGLBuffers;
    
	public JoclGpuCompute(GLAutoDrawable drawable, boolean... verbose) {
		//set debug property
		if(verbose.length > 0)
			this.verbose = verbose[0];
		
		// find gl compatible device
        CLDevice[] devices = CLPlatform.getDefault().listCLDevices();
        CLDevice device = null;
        for (CLDevice d : devices) {
            if(d.isGLMemorySharingSupported()) {
                device = d;
                break;
            }
        }
        if(null==device) {
            throw new RuntimeException("JoclGpuCompute: couldn't find any CL/GL memory sharing devices ..");
        }
        // create OpenCL context before creating any OpenGL objects
        // you want to share with OpenCL (AMD driver requirement)
        clGLContext = CLGLContext.create(drawable.getContext(), device);
        usingGLInterop = true;
	}
	
	public JoclGpuCompute(int profile, boolean... debug){        
		CLPlatform platform = CLPlatform.getDefault();
        
		if(profile == CROWD_OCL_CPU_ONLY){
			clContext = CLContext.create(platform, Type.CPU);
		}else if(profile == CROWD_OCL_GPU_ONLY){
			clContext = CLContext.create(platform, Type.GPU);
		}else if(profile == CROWD_OCL_BEST_AVAILABLE){
			clContext = CLContext.create(platform.getMaxFlopsDevice());
		}else{
			throw new RuntimeException("JoclGpuCompute: bad profile index");
		}
    }
	
	public void initCL(String programSource){
		CLProgram program = compileCLCode(programSource);
		initKernelMap(program);
		initBuffers();
	}
	
	public void initCL(InputStream programFile) {
        CLProgram program = compileCLCode(programFile);
        initKernelMap(program);
        initBuffers();
    }
	
	public void addKernelFile(String programSource){
		CLProgram program = compileCLCode(programSource);
		HashMap<String, CLKernel> temp = (HashMap<String, CLKernel>)program.createCLKernels();
		kernels.putAll(temp);
		if(verbose)
			System.out.println(kernels.keySet());
	}
	
	public void addKernelFile(InputStream programFile){
		CLProgram program = compileCLCode(programFile);
		HashMap<String, CLKernel> temp = (HashMap<String, CLKernel>)program.createCLKernels();
		kernels.putAll(temp);
		if(verbose)
			System.out.println(kernels.keySet());
	}
	
	private void initKernelMap(CLProgram program){
		kernels = (HashMap<String, CLKernel>) program.createCLKernels();
        if(verbose)
        	System.out.println(kernels.keySet());
	}
	
	private void initBuffers(){
		clBuffers = new HashMap<String, CLBuffer<FloatBuffer>>();
        if(usingGLInterop){
        	commandQueue = clGLContext.getMaxFlopsDevice().createCommandQueue();
        	clGLBuffers = new HashMap<String, CLGLBuffer<?>>();
        }else{
        	commandQueue = clContext.getMaxFlopsDevice().createCommandQueue();
        }
        if(verbose)
        	System.out.println("cl initialised");
	}
	
	private CLProgram compileCLCode(String programSource){
		CLProgram program;
        if(usingGLInterop)
        	program = clGLContext.createProgram(programSource);
        else
        	program = clContext.createProgram(programSource);
        program.build();
        
        if(verbose){
        	System.out.println(program.getBuildStatus());
        	System.out.println(program.isExecutable());
        	System.out.println(program.getBuildLog());
        }
        
        return program;
	}
	
	private CLProgram compileCLCode(InputStream programFile){
		CLProgram program;
        try {
        	if(usingGLInterop)
        		program = clGLContext.createProgram(programFile);
        	else
        		program = clContext.createProgram(programFile);
            program.build();
            if(verbose){
            	System.out.println(program.getBuildStatus());
            	System.out.println(program.isExecutable());
            	System.out.println(program.getBuildLog());
            }
        } catch (IOException ex) {
            throw new RuntimeException("JoclGpuCompute: can not build program", ex);
        }
        return program;
	}
	
	public void setKernelArgs(String kernelName, CLMemory<?>... values){
		kernels.get(kernelName).putArgs(values).rewind();
	}
	
	public void setKernelArg(String kernelName, int index, CLMemory<?> value){
		kernels.get(kernelName).setArg(index, value);
	}
	
	public void setKernelArg(String kernelName, int index, float value){
		kernels.get(kernelName).setArg(index, value);
	}
	
	public void setKernelArg(String kernelName, int index, int value){
		kernels.get(kernelName).setArg(index, value);
	}
	
	public void createCLBuffer(String name, int size, CLMemory.Mem... flags){
		clBuffers.put(name, clContext.createFloatBuffer(size, flags));
		//CL c = clContext.getCL();
		//c.clCreateBuffer(clContext.ID, null, (long)size, Buffers.newDirectFloatBuffer(size)., null);
		//clBuffers.put(name, createBuffer(newDirectFloatBuffer(size), flags));
	}
    
    /**
     * Creates a CLBuffer with the specified flags. No flags creates a MEM.READ_WRITE buffer.
     */
    /*public final <B extends Buffer> CLBuffer<B> createBuffer(B directBuffer, Mem... flags) {
        return createBuffer(directBuffer, Mem.flagsToInt(flags));
    }*/

    /**
     * Creates a CLBuffer with the specified flags.
     */
    /*public final <B extends Buffer> CLBuffer<B> createBuffer(B directBuffer, int flags) {
        CLBuffer<B> buffer = CLBuffer.create(clContext, directBuffer, flags);
        //memoryObjects.add(buffer);
        return buffer;
    }*/
	
	public void releaseCLBuffer(String name){
		CL c = clContext.getCL();
		c.clReleaseMemObject(clBuffers.get(name).ID);
	}
	
	public void fillBuffer(String name, float src[]){
		FloatBuffer buffer = (FloatBuffer) clBuffers.get(name).getBuffer();
		if(buffer != null && src != null){
			buffer.put(src);
		}
		buffer.rewind();
		commandQueue.putWriteBuffer(clBuffers.get(name), false).finish();
	}
	
	public void fillBuffer(String name, Float src[]){
		FloatBuffer buffer = (FloatBuffer) clBuffers.get(name).getBuffer();
		if(buffer != null && src != null){
			for(int i=0; i<src.length; i++)
				buffer.put(src[i]);
		}
		buffer.rewind();
		commandQueue.putWriteBuffer(clBuffers.get(name), false).finish();
	}
	
	public void createCLGLBuffer(GLAutoDrawable drawable, String name, int size, CLMemory.Mem... clFlags){
		final int[] id = new int[1];
		GL2 gl = drawable.getGL().getGL2();
		gl.glGenBuffers(1, id, 0);
		gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, id[0]);
        gl.glBufferData(GL2.GL_ARRAY_BUFFER, size, null, GL2.GL_DYNAMIC_DRAW);
        gl.glBindBuffer(GL2.GL_ARRAY_BUFFER, 0);
        
        clGLBuffers.put(name, clGLContext.createFromGLBuffer(id[0], clFlags));
        
        if(verbose){
        	System.out.println("cl buffer type: " + clGLBuffers.get(name).getGLObjectType());
        	System.out.println("shared with gl buffer: " + clGLBuffers.get(name).getGLObjectID());
        }
	}
	
	public void createCLGLBuffer(String name, int size, CLMemory.Mem... clFlags){
		clBuffers.put(name, clGLContext.createFloatBuffer(size, clFlags));
	}
	
	public void releaseCLGLBuffer(String name){
		CL c = clGLContext.getCL();
		c.clReleaseMemObject(clBuffers.get(name).ID);
	}
	
	
	
	public void compute() {
	}
	
	public int getVertexObjectId(String buffer) {
		return this.clGLBuffers.get(buffer).getGLObjectID();
	}
	
	public int getArraySize(String buffer) {
		return 0;//MESH_SIZE * MESH_SIZE;
	}
	
}
